import os, csv, torch, numpy, scipy.io, PIL.Image, torchvision.transforms
from mit_utils import colorEncode
from model import EncoderDecoder
import jittor as jt

jt.flags.use_cuda = 1

# TODO: Image path
img_path = './demo/ADE_val_00000003.jpg'

colors = scipy.io.loadmat('./data/color150.mat')['colors']

names = {}

with open('./data/object150_info.csv') as f:
    reader = csv.reader(f)
    next(reader)
    for row in reader:
        names[int(row[0])] = row[5].split(";")[0]

def visualize_result(pred, index=None):
    # filter prediction class if requested
    if index is not None:
        pred = pred.copy()
        pred[pred != index] = -1
        print(f'{names[index+1]}:')
        
    # colorize prediction
    pred_color = colorEncode(pred, colors).astype(numpy.uint8)

    # aggregate images and save
    im_vis = pred_color
    PIL.Image.fromarray(im_vis).save(
        os.path.join(img_path.replace('.jpg', '.png')))


# Network
# TODO: change to your checkpoint path
resume = "./ckpt/ade20k-resnet101-new-cca_deepsup/epoch_40.pkl"
segmentation_module = EncoderDecoder(resume=resume)
segmentation_module.eval()

# load image
pil_to_tensor = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize(
        mean=[0.485, 0.456, 0.406], # These are RGB mean+std values
        std=[0.229, 0.224, 0.225])  # across a large photo dataset.
])
pil_image = PIL.Image.open(img_path).convert('RGB')
img_original = numpy.array(pil_image)
img_data = pil_to_tensor(pil_image)
singleton_batch = {'img_data': img_data[None], 'seg_label': torch.tensor([1])} # ignore seg_label please
output_size = img_data.shape[1:]

# Run the segmentation
scores = segmentation_module(singleton_batch, segSize=output_size)

# Get the predicted scores for each pixel
pred = jt.argmax(scores, dim=1)[0][0].numpy()
visualize_result(pred)
